{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "89d7e754-8d3b-450e-8ca2-e21f91aecdd3",
   "metadata": {},
   "source": [
    "> (BN) BatchNorm\\\n",
    "> (LN) LayerNorm\\\n",
    "> (RMS) RMSNorm\n",
    "\n",
    "- it has been the standard to use\n",
    "    - batchnorm in CV tasks,\n",
    "    - layernorm in NLP tasks.\n",
    "- references\n",
    "    - https://afterhoursresearch.hashnode.dev/batch-normalization-layer-normalization-and-root-mean-square-layer-normalization-a-comprehensive-guide-with-python-implementations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8fbcd49-8641-4d48-b1bf-18addb54140b",
   "metadata": {},
   "source": [
    "- basics\n",
    "    - dnn 的标准组件，稳定和加速训练过程\n",
    "    - `batch_size, seq_len, hidden_size = 2, 3, 4`\n",
    "    - Batch Norm: reduce cross batch size\n",
    "        - mini-batch dimension\n",
    "        - 一般用于图像，不涉及到padding的问题；\n",
    "    - Layer Norm: reduce cross hidden dim\n",
    "        - 一般用于序列，一个 batch size 内存在 padding；\n",
    "        - across the feature dimension.\n",
    "        - RMSNorm: 对 LN 的一种变体，llama\n",
    "    - residual connection + norm\n",
    "        - https://spaces.ac.cn/archives/9009\n",
    "        - Pre LN: `llama`\n",
    "        - Post LN: `attention is all you need`\n",
    "- 理解高维 tensor 的 shape （axis/dim）处理\n",
    "    - 最终的输出都不改变 shape\n",
    "- 理解这三个 norm 的计算过程"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "65c0088e-1263-441d-a5e4-997c79edb04b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-16T03:20:57.289320Z",
     "iopub.status.busy": "2024-09-16T03:20:57.288752Z",
     "iopub.status.idle": "2024-09-16T03:20:58.868689Z",
     "shell.execute_reply": "2024-09-16T03:20:58.867917Z",
     "shell.execute_reply.started": "2024-09-16T03:20:57.289278Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch._C.Generator at 0x736c404e39f0>"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "import transformers\n",
    "from transformers import AutoModelForCausalLM\n",
    "\n",
    "from IPython.display import Image\n",
    "import os\n",
    "os.environ['http_proxy'] = 'http://127.0.0.1:7890'\n",
    "os.environ['https_proxy'] = 'http://127.0.0.1:7890'\n",
    "\n",
    "torch.manual_seed(42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0d1b8a96-440a-4767-8138-a1403a500c9f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-16T03:21:00.773392Z",
     "iopub.status.busy": "2024-09-16T03:21:00.772845Z",
     "iopub.status.idle": "2024-09-16T03:21:00.778172Z",
     "shell.execute_reply": "2024-09-16T03:21:00.777282Z",
     "shell.execute_reply.started": "2024-09-16T03:21:00.773371Z"
    }
   },
   "outputs": [],
   "source": [
    "# 3d tensor\n",
    "batch_size, seq_len, hidden_size = 2, 3, 4\n",
    "x = torch.randn(batch_size, seq_len, hidden_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6c1acbe1-a8f0-467a-bee8-b257d4198ffd",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-16T03:21:02.443503Z",
     "iopub.status.busy": "2024-09-16T03:21:02.442527Z",
     "iopub.status.idle": "2024-09-16T03:21:02.455547Z",
     "shell.execute_reply": "2024-09-16T03:21:02.454287Z",
     "shell.execute_reply.started": "2024-09-16T03:21:02.443447Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 1.9269,  1.4873,  0.9007, -2.1055],\n",
       "         [ 0.6784, -1.2345, -0.0431, -1.6047],\n",
       "         [ 0.3559, -0.6866, -0.4934,  0.2415]],\n",
       "\n",
       "        [[-1.1109,  0.0915, -2.3169, -0.2168],\n",
       "         [-0.3097, -0.3957,  0.8034, -0.6216],\n",
       "         [-0.5920, -0.0631, -0.8286,  0.3309]]])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b51a3a68-a60a-4d58-8143-1c0ade7b9021",
   "metadata": {},
   "source": [
    "### BN"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d04a833-9029-4537-8270-b1e06745269b",
   "metadata": {},
   "source": [
    "x: (a, b, c)\n",
    "- mean/var 都是 reduce 的过程；\n",
    "- mean(dim=a) => 返回的 shape 是 (b, c)\n",
    "    - mean(dim=a, keepdim=True) => (1, b, c): 方便 broadcast\n",
    "    - einsum: `ijk->jk`\n",
    "- mean(dim=(a, b)) => 返回的 shape 是 (c)\n",
    "    - mean(dim=(a, b), keepdim=True) => (1, 1, c): 方便 broadcast\n",
    "    - einsum: `ijk->k`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d89e3025-fcb6-4765-b85d-9e3ed13d0b9d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-16T03:21:05.399247Z",
     "iopub.status.busy": "2024-09-16T03:21:05.397452Z",
     "iopub.status.idle": "2024-09-16T03:21:05.406385Z",
     "shell.execute_reply": "2024-09-16T03:21:05.404997Z",
     "shell.execute_reply.started": "2024-09-16T03:21:05.399179Z"
    }
   },
   "outputs": [],
   "source": [
    "bn = nn.BatchNorm1d(hidden_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9c696211-2c82-4e93-b8d7-cfeea17a48b8",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-16T03:21:10.370701Z",
     "iopub.status.busy": "2024-09-16T03:21:10.370047Z",
     "iopub.status.idle": "2024-09-16T03:21:11.881644Z",
     "shell.execute_reply": "2024-09-16T03:21:11.879295Z",
     "shell.execute_reply.started": "2024-09-16T03:21:10.370653Z"
    }
   },
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "running_mean should contain 3 elements not 4",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[5], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mbn\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/anaconda3/lib/python3.10/site-packages/torch/nn/modules/module.py:1553\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1551\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1552\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1553\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/anaconda3/lib/python3.10/site-packages/torch/nn/modules/module.py:1562\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1557\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1558\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1559\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1560\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1561\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1562\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1564\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1565\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "File \u001b[0;32m~/anaconda3/lib/python3.10/site-packages/torch/nn/modules/batchnorm.py:176\u001b[0m, in \u001b[0;36m_BatchNorm.forward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    169\u001b[0m     bn_training \u001b[38;5;241m=\u001b[39m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrunning_mean \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m) \u001b[38;5;129;01mand\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrunning_var \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m)\n\u001b[1;32m    171\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124mr\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    172\u001b[0m \u001b[38;5;124;03mBuffers are only updated if they are to be tracked and we are in training mode. Thus they only need to be\u001b[39;00m\n\u001b[1;32m    173\u001b[0m \u001b[38;5;124;03mpassed when the update should occur (i.e. in training mode when they are tracked), or when buffer stats are\u001b[39;00m\n\u001b[1;32m    174\u001b[0m \u001b[38;5;124;03mused for normalization (i.e. in eval mode when buffers are not None).\u001b[39;00m\n\u001b[1;32m    175\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[0;32m--> 176\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mF\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbatch_norm\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    177\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m    178\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;66;43;03m# If buffers are not to be tracked, ensure that they won't be updated\u001b[39;49;00m\n\u001b[1;32m    179\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrunning_mean\u001b[49m\n\u001b[1;32m    180\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43;01mif\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mnot\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtraining\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtrack_running_stats\u001b[49m\n\u001b[1;32m    181\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43;01melse\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m    182\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrunning_var\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mif\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mnot\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtraining\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtrack_running_stats\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01melse\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m    183\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mweight\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    184\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbias\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    185\u001b[0m \u001b[43m    \u001b[49m\u001b[43mbn_training\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    186\u001b[0m \u001b[43m    \u001b[49m\u001b[43mexponential_average_factor\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    187\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43meps\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    188\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/anaconda3/lib/python3.10/site-packages/torch/nn/functional.py:2512\u001b[0m, in \u001b[0;36mbatch_norm\u001b[0;34m(input, running_mean, running_var, weight, bias, training, momentum, eps)\u001b[0m\n\u001b[1;32m   2509\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m training:\n\u001b[1;32m   2510\u001b[0m     _verify_batch_size(\u001b[38;5;28minput\u001b[39m\u001b[38;5;241m.\u001b[39msize())\n\u001b[0;32m-> 2512\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbatch_norm\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   2513\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mweight\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbias\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mrunning_mean\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mrunning_var\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtraining\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmomentum\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43meps\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbackends\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcudnn\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43menabled\u001b[49m\n\u001b[1;32m   2514\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: running_mean should contain 3 elements not 4"
     ]
    }
   ],
   "source": [
    "bn(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a28d60a7-27bf-4a14-8743-55c3c9b179f2",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-16T02:51:57.671699Z",
     "iopub.status.busy": "2024-09-16T02:51:57.671499Z",
     "iopub.status.idle": "2024-09-16T02:51:57.683628Z",
     "shell.execute_reply": "2024-09-16T02:51:57.682970Z",
     "shell.execute_reply.started": "2024-09-16T02:51:57.671684Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 1.7943,  1.9214,  1.1306, -1.5846],\n",
       "         [ 0.5278, -1.3052,  0.2633, -1.0345],\n",
       "         [ 0.2006, -0.6557, -0.1505,  0.9930]],\n",
       "\n",
       "        [[-1.2873,  0.2668, -1.8263,  0.4897],\n",
       "         [-0.4746, -0.3108,  1.0412,  0.0451],\n",
       "         [-0.7609,  0.0835, -0.4585,  1.0912]]], grad_fn=<TransposeBackward0>)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bn(x.transpose(1, 2)).transpose(1, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8be76a65-5efc-4639-b968-a1022fb10edb",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-16T02:51:57.686526Z",
     "iopub.status.busy": "2024-09-16T02:51:57.686216Z",
     "iopub.status.idle": "2024-09-16T02:51:57.692778Z",
     "shell.execute_reply": "2024-09-16T02:51:57.691821Z",
     "shell.execute_reply.started": "2024-09-16T02:51:57.686509Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 1.9269,  0.6784,  0.3559],\n",
       "         [ 1.4873, -1.2345, -0.6866],\n",
       "         [ 0.9007, -0.0431, -0.4934],\n",
       "         [-2.1055, -1.6047,  0.2415]],\n",
       "\n",
       "        [[-1.1109, -0.3097, -0.5920],\n",
       "         [ 0.0915, -0.3957, -0.0631],\n",
       "         [-2.3169,  0.8034, -0.8286],\n",
       "         [-0.2168, -0.6216,  0.3309]]])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 2*3*4 => 2*4*3\n",
    "x.transpose(1, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9f877a4d-4a74-42d9-b717-4493aca1b293",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-16T02:51:57.693557Z",
     "iopub.status.busy": "2024-09-16T02:51:57.693363Z",
     "iopub.status.idle": "2024-09-16T02:51:57.706214Z",
     "shell.execute_reply": "2024-09-16T02:51:57.704935Z",
     "shell.execute_reply.started": "2024-09-16T02:51:57.693543Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 0.1581],\n",
       "         [-0.1335],\n",
       "         [-0.3296],\n",
       "         [-0.6627]]])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mean = x.transpose(1, 2).mean(dim=(0, 2), keepdim=True)\n",
    "mean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "4ca4ea21-942c-4f8f-9b2d-ea64d55b37b4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-16T03:16:37.588994Z",
     "iopub.status.busy": "2024-09-16T03:16:37.588383Z",
     "iopub.status.idle": "2024-09-16T03:16:37.603698Z",
     "shell.execute_reply": "2024-09-16T03:16:37.601730Z",
     "shell.execute_reply.started": "2024-09-16T03:16:37.588950Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 0.1581, -0.1335, -0.3296, -0.6627])"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 序列的每个位置算均值\n",
    "torch.einsum('ijk->k', x) / (x.shape[0] * x.shape[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f7053feb-1515-438b-83e7-29695b870b05",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-16T03:49:30.289482Z",
     "iopub.status.busy": "2024-09-16T03:49:30.288844Z",
     "iopub.status.idle": "2024-09-16T03:49:30.304098Z",
     "shell.execute_reply": "2024-09-16T03:49:30.302004Z",
     "shell.execute_reply.started": "2024-09-16T03:49:30.289435Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 0.1581],\n",
       "         [-0.1335],\n",
       "         [-0.3296],\n",
       "         [-0.6627]]])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# x.transpose(1, 2).mean(dim=0).mean(dim=1)\n",
    "x.transpose(1, 2).mean(dim=0, keepdim=True).mean(dim=2, keepdim=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "dae1ce93-d493-47a3-a6af-d39163f984ea",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-16T03:01:54.907242Z",
     "iopub.status.busy": "2024-09-16T03:01:54.906634Z",
     "iopub.status.idle": "2024-09-16T03:01:54.922820Z",
     "shell.execute_reply": "2024-09-16T03:01:54.920802Z",
     "shell.execute_reply.started": "2024-09-16T03:01:54.907198Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 1.9269,  0.6784,  0.3559, -1.1109, -0.3097, -0.5920],\n",
       "        [ 1.4873, -1.2345, -0.6866,  0.0915, -0.3957, -0.0631],\n",
       "        [ 0.9007, -0.0431, -0.4934, -2.3169,  0.8034, -0.8286],\n",
       "        [-2.1055, -1.6047,  0.2415, -0.2168, -0.6216,  0.3309]])"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 2*4*3 => 4*6\n",
    "x.transpose(1, 2).transpose(0, 1).reshape(4, -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "fde34430-b62c-454f-beb0-449f61a4a3ca",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-16T03:02:11.551311Z",
     "iopub.status.busy": "2024-09-16T03:02:11.550707Z",
     "iopub.status.idle": "2024-09-16T03:02:11.564531Z",
     "shell.execute_reply": "2024-09-16T03:02:11.562524Z",
     "shell.execute_reply.started": "2024-09-16T03:02:11.551267Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 0.1581, -0.1335, -0.3296, -0.6627])"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.transpose(1, 2).transpose(0, 1).reshape(4, -1).mean(dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "887fd4dd-f401-442d-b85c-93c2f80b9f35",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-16T02:51:57.707583Z",
     "iopub.status.busy": "2024-09-16T02:51:57.707378Z",
     "iopub.status.idle": "2024-09-16T02:51:57.714659Z",
     "shell.execute_reply": "2024-09-16T02:51:57.713790Z",
     "shell.execute_reply.started": "2024-09-16T02:51:57.707568Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[0.9717],\n",
       "         [0.7116],\n",
       "         [1.1841],\n",
       "         [0.8291]]])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "var = x.transpose(1, 2).var(dim=(0, 2), keepdim=True, unbiased=False)\n",
    "var"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ea109c4c-34c0-4274-a40c-c81886439f50",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-16T02:51:57.715588Z",
     "iopub.status.busy": "2024-09-16T02:51:57.715388Z",
     "iopub.status.idle": "2024-09-16T02:51:57.723077Z",
     "shell.execute_reply": "2024-09-16T02:51:57.722226Z",
     "shell.execute_reply.started": "2024-09-16T02:51:57.715573Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 1.7943,  1.9214,  1.1306, -1.5846],\n",
       "         [ 0.5278, -1.3052,  0.2633, -1.0345],\n",
       "         [ 0.2006, -0.6557, -0.1505,  0.9930]],\n",
       "\n",
       "        [[-1.2873,  0.2668, -1.8263,  0.4897],\n",
       "         [-0.4746, -0.3108,  1.0412,  0.0451],\n",
       "         [-0.7609,  0.0835, -0.4585,  1.0912]]])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "((x.transpose(1, 2) - mean) / torch.sqrt(var + bn.eps)).transpose(1, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "dd7c2414-4ca9-4688-b5a5-4646bc1d66a8",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-16T02:51:57.724186Z",
     "iopub.status.busy": "2024-09-16T02:51:57.723984Z",
     "iopub.status.idle": "2024-09-16T02:51:57.731009Z",
     "shell.execute_reply": "2024-09-16T02:51:57.730047Z",
     "shell.execute_reply.started": "2024-09-16T02:51:57.724171Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 1.7943,  1.9214,  1.1306, -1.5846],\n",
       "         [ 0.5278, -1.3052,  0.2633, -1.0345],\n",
       "         [ 0.2006, -0.6557, -0.1505,  0.9930]],\n",
       "\n",
       "        [[-1.2873,  0.2668, -1.8263,  0.4897],\n",
       "         [-0.4746, -0.3108,  1.0412,  0.0451],\n",
       "         [-0.7609,  0.0835, -0.4585,  1.0912]]])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(x - x.mean(dim=(0, 1), keepdim=True)) / torch.sqrt(x.var(dim=(0, 1), unbiased=False, keepdim=True) + bn.eps)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1226d86e-6af3-49f7-88cf-d5bd8a4d23ee",
   "metadata": {},
   "source": [
    "### LN"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7416136b-d91f-4ce7-af67-afb2773a664f",
   "metadata": {},
   "source": [
    "- element-wise operation\n",
    "\n",
    "    $$\n",
    "    y = \\frac{x - \\mathbb{E}[x]}{\\sqrt{\\mathrm{Var}[x] + \\epsilon}} * \\gamma + \\beta\n",
    "    $$\n",
    "\n",
    "- 关于 $\\mathrm{Var}[x]$\n",
    "  \n",
    "    $$\n",
    "    \\mathrm{Var}[x] = \\mathbb{E}[(x - \\mathbb{E}[x])^2]\n",
    "    $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "848ba484-4eb7-4191-a067-a17943737680",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-16T02:51:57.732747Z",
     "iopub.status.busy": "2024-09-16T02:51:57.732212Z",
     "iopub.status.idle": "2024-09-16T02:51:57.739695Z",
     "shell.execute_reply": "2024-09-16T02:51:57.739005Z",
     "shell.execute_reply.started": "2024-09-16T02:51:57.732713Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 0.8716,  0.5928,  0.2209, -1.6853],\n",
       "         [ 1.3440, -0.7473,  0.5552, -1.1519],\n",
       "         [ 1.1111, -1.1985, -0.7703,  0.8577]],\n",
       "\n",
       "        [[-0.2380,  1.0472, -1.5270,  0.7177],\n",
       "         [-0.3243, -0.4803,  1.6947, -0.8900],\n",
       "         [-0.6717,  0.4977, -1.1947,  1.3688]]],\n",
       "       grad_fn=<NativeLayerNormBackward0>)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ln = nn.LayerNorm(hidden_size)\n",
    "ln_out = ln(x)\n",
    "ln_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "cc8ea066-28a0-41d9-bc9d-bccb1c6f93a1",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-16T02:51:57.741376Z",
     "iopub.status.busy": "2024-09-16T02:51:57.740754Z",
     "iopub.status.idle": "2024-09-16T02:51:57.746807Z",
     "shell.execute_reply": "2024-09-16T02:51:57.745978Z",
     "shell.execute_reply.started": "2024-09-16T02:51:57.741342Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "LayerNorm((4,), eps=1e-05, elementwise_affine=True)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ln"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "efb8809b-33ac-43c8-b132-5832ec5813fe",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-16T02:51:57.747875Z",
     "iopub.status.busy": "2024-09-16T02:51:57.747682Z",
     "iopub.status.idle": "2024-09-16T02:51:57.755696Z",
     "shell.execute_reply": "2024-09-16T02:51:57.754851Z",
     "shell.execute_reply.started": "2024-09-16T02:51:57.747861Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 1.9269,  1.4873,  0.9007, -2.1055])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x[0, 0, :]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "6984a30c-0f6e-45bb-918d-e92f171110af",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-16T02:51:57.756764Z",
     "iopub.status.busy": "2024-09-16T02:51:57.756569Z",
     "iopub.status.idle": "2024-09-16T02:51:57.764519Z",
     "shell.execute_reply": "2024-09-16T02:51:57.763673Z",
     "shell.execute_reply.started": "2024-09-16T02:51:57.756750Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 0.8716,  0.5928,  0.2209, -1.6853])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(x[0, 0, :] - torch.mean(x[0, 0, :])) / torch.sqrt(torch.var(x[0, 0, :], unbiased=False) + ln.eps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "4b14e7af-e170-4de2-8650-21e28a4cbeef",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-16T02:51:57.765602Z",
     "iopub.status.busy": "2024-09-16T02:51:57.765399Z",
     "iopub.status.idle": "2024-09-16T02:51:57.772742Z",
     "shell.execute_reply": "2024-09-16T02:51:57.771907Z",
     "shell.execute_reply.started": "2024-09-16T02:51:57.765587Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 0.8716,  0.5928,  0.2209, -1.6853],\n",
       "         [ 1.3440, -0.7473,  0.5552, -1.1519],\n",
       "         [ 1.1111, -1.1985, -0.7703,  0.8577]],\n",
       "\n",
       "        [[-0.2380,  1.0472, -1.5270,  0.7177],\n",
       "         [-0.3243, -0.4803,  1.6947, -0.8900],\n",
       "         [-0.6717,  0.4977, -1.1947,  1.3688]]])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(x - x.mean(dim=-1, keepdim=True)) / torch.sqrt(x.var(dim=-1, unbiased=False, keepdim=True) + ln.eps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "4faf3098-b5dd-437b-b2fa-3b2235807628",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-16T03:17:14.849209Z",
     "iopub.status.busy": "2024-09-16T03:17:14.848543Z",
     "iopub.status.idle": "2024-09-16T03:17:14.864099Z",
     "shell.execute_reply": "2024-09-16T03:17:14.861871Z",
     "shell.execute_reply.started": "2024-09-16T03:17:14.849162Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 1.3746,  0.9349,  0.3484, -2.6579],\n",
       "         [ 1.2294, -0.6836,  0.5079, -1.0537],\n",
       "         [ 0.5015, -0.5410, -0.3477,  0.3871]],\n",
       "\n",
       "        [[-0.2226,  0.9798, -1.4287,  0.6715],\n",
       "         [-0.1788, -0.2648,  0.9343, -0.4907],\n",
       "         [-0.3038,  0.2251, -0.5404,  0.6191]]])"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(x - x.mean(dim=-1, keepdim=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "28c34985-831c-4000-8247-f170dfd3f07e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-16T03:19:23.102203Z",
     "iopub.status.busy": "2024-09-16T03:19:23.101496Z",
     "iopub.status.idle": "2024-09-16T03:19:23.118488Z",
     "shell.execute_reply": "2024-09-16T03:19:23.116363Z",
     "shell.execute_reply.started": "2024-09-16T03:19:23.102153Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 1.3746,  0.9349,  0.3484, -2.6579],\n",
       "         [ 1.2294, -0.6836,  0.5079, -1.0537],\n",
       "         [ 0.5015, -0.5410, -0.3477,  0.3871]],\n",
       "\n",
       "        [[-0.2226,  0.9798, -1.4287,  0.6715],\n",
       "         [-0.1788, -0.2648,  0.9343, -0.4907],\n",
       "         [-0.3038,  0.2251, -0.5404,  0.6191]]])"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x - torch.einsum('ijk->ij', x).unsqueeze(2) / x.shape[-1]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fed12c0-fb7b-42ee-91a0-ef658a091aaf",
   "metadata": {},
   "source": [
    "### RMSNorm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "641afccf-b8e9-47e2-a7f1-671e1f1467bd",
   "metadata": {},
   "source": [
    "\n",
    "$$\n",
    "y = \\frac{x}{\\sqrt{\\frac1n\\sum_ix_i^2+\\epsilon}} * \\gamma\n",
    "$$\n",
    "\n",
    "- https://github.com/meta-llama/llama3/blob/main/llama/model.py#L35C7-L46\n",
    "- 存在 learnable parameter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "ccfea7bb-99d5-4d2b-ae80-ebb56a2ed149",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-16T02:51:57.773823Z",
     "iopub.status.busy": "2024-09-16T02:51:57.773614Z",
     "iopub.status.idle": "2024-09-16T02:51:57.780019Z",
     "shell.execute_reply": "2024-09-16T02:51:57.779130Z",
     "shell.execute_reply.started": "2024-09-16T02:51:57.773807Z"
    }
   },
   "outputs": [],
   "source": [
    "class RMSNorm(torch.nn.Module):\n",
    "    def __init__(self, dim: int, eps: float = 1e-6):\n",
    "        super().__init__()\n",
    "        self.eps = eps\n",
    "        self.weight = nn.Parameter(torch.ones(dim))\n",
    "\n",
    "    def _norm(self, x):\n",
    "        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)\n",
    "\n",
    "    def forward(self, x):\n",
    "        output = self._norm(x.float()).type_as(x)\n",
    "        return output * self.weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "ee908e9f-d840-45fa-9f7a-03f304d8fb20",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-16T02:51:57.781317Z",
     "iopub.status.busy": "2024-09-16T02:51:57.781117Z",
     "iopub.status.idle": "2024-09-16T02:51:57.790753Z",
     "shell.execute_reply": "2024-09-16T02:51:57.789905Z",
     "shell.execute_reply.started": "2024-09-16T02:51:57.781301Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 1.1531,  0.8900,  0.5390, -1.2600],\n",
       "         [ 0.6353, -1.1561, -0.0403, -1.5027],\n",
       "         [ 0.7503, -1.4477, -1.0402,  0.5092]],\n",
       "\n",
       "        [[-0.8611,  0.0710, -1.7959, -0.1681],\n",
       "         [-0.5466, -0.6983,  1.4178, -1.0970],\n",
       "         [-1.1039, -0.1176, -1.5450,  0.6170]]], grad_fn=<MulBackward0>)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rms = RMSNorm(dim=4)\n",
    "rms(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "b5d1fdc5-1335-4da1-aca6-9ec9f0312d1b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-16T02:51:57.791907Z",
     "iopub.status.busy": "2024-09-16T02:51:57.791687Z",
     "iopub.status.idle": "2024-09-16T02:51:57.799012Z",
     "shell.execute_reply": "2024-09-16T02:51:57.798088Z",
     "shell.execute_reply.started": "2024-09-16T02:51:57.791890Z"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 1.1531,  0.8900,  0.5390, -1.2600],\n",
       "         [ 0.6353, -1.1561, -0.0403, -1.5027],\n",
       "         [ 0.7503, -1.4477, -1.0402,  0.5092]],\n",
       "\n",
       "        [[-0.8611,  0.0710, -1.7959, -0.1681],\n",
       "         [-0.5466, -0.6983,  1.4178, -1.0970],\n",
       "         [-1.1039, -0.1176, -1.5450,  0.6170]]])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x / torch.sqrt(x.pow(2).mean(dim=-1, keepdim=True) + 1e-6)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a8a039f-eeb6-4976-8f15-bf1615e1d962",
   "metadata": {},
   "source": [
    "#### llama3 rmsnorm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "0307ce84-6317-49b3-8c5c-502753e85923",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-16T02:51:57.799902Z",
     "iopub.status.busy": "2024-09-16T02:51:57.799672Z",
     "iopub.status.idle": "2024-09-16T02:52:03.892244Z",
     "shell.execute_reply": "2024-09-16T02:52:03.891221Z",
     "shell.execute_reply.started": "2024-09-16T02:51:57.799885Z"
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3b7a6a0e20144a6c85128681146c3f1a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "We've detected an older driver with an RTX 4000 series GPU. These drivers have issues with P2P. This can affect the multi-gpu inference when using accelerate device_map.Please make sure to update your driver to the latest version which resolves this.\n"
     ]
    }
   ],
   "source": [
    "llama3_id = \"meta-llama/Meta-Llama-3-8B\"\n",
    "llama3 = AutoModelForCausalLM.from_pretrained(llama3_id, torch_dtype=torch.bfloat16, device_map='auto')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "068a3859-39df-43cc-858a-691ebae74e61",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-16T02:52:03.893662Z",
     "iopub.status.busy": "2024-09-16T02:52:03.893329Z",
     "iopub.status.idle": "2024-09-16T02:52:03.901993Z",
     "shell.execute_reply": "2024-09-16T02:52:03.901080Z",
     "shell.execute_reply.started": "2024-09-16T02:52:03.893643Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "LlamaForCausalLM(\n",
       "  (model): LlamaModel(\n",
       "    (embed_tokens): Embedding(128256, 4096)\n",
       "    (layers): ModuleList(\n",
       "      (0-31): 32 x LlamaDecoderLayer(\n",
       "        (self_attn): LlamaSdpaAttention(\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=1024, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=1024, bias=False)\n",
       "          (o_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (rotary_emb): LlamaRotaryEmbedding()\n",
       "        )\n",
       "        (mlp): LlamaMLP(\n",
       "          (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)\n",
       "          (up_proj): Linear(in_features=4096, out_features=14336, bias=False)\n",
       "          (down_proj): Linear(in_features=14336, out_features=4096, bias=False)\n",
       "          (act_fn): SiLU()\n",
       "        )\n",
       "        (input_layernorm): LlamaRMSNorm((4096,), eps=1e-05)\n",
       "        (post_attention_layernorm): LlamaRMSNorm((4096,), eps=1e-05)\n",
       "      )\n",
       "    )\n",
       "    (norm): LlamaRMSNorm((4096,), eps=1e-05)\n",
       "    (rotary_emb): LlamaRotaryEmbedding()\n",
       "  )\n",
       "  (lm_head): Linear(in_features=4096, out_features=128256, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "llama3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "54326702-e70b-4272-b10d-ce265d94ab59",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-16T02:52:03.905415Z",
     "iopub.status.busy": "2024-09-16T02:52:03.904905Z",
     "iopub.status.idle": "2024-09-16T02:52:03.921933Z",
     "shell.execute_reply": "2024-09-16T02:52:03.921036Z",
     "shell.execute_reply.started": "2024-09-16T02:52:03.905395Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "LlamaConfig {\n",
       "  \"_name_or_path\": \"meta-llama/Meta-Llama-3-8B\",\n",
       "  \"architectures\": [\n",
       "    \"LlamaForCausalLM\"\n",
       "  ],\n",
       "  \"attention_bias\": false,\n",
       "  \"attention_dropout\": 0.0,\n",
       "  \"bos_token_id\": 128000,\n",
       "  \"eos_token_id\": 128001,\n",
       "  \"head_dim\": 128,\n",
       "  \"hidden_act\": \"silu\",\n",
       "  \"hidden_size\": 4096,\n",
       "  \"initializer_range\": 0.02,\n",
       "  \"intermediate_size\": 14336,\n",
       "  \"max_position_embeddings\": 8192,\n",
       "  \"mlp_bias\": false,\n",
       "  \"model_type\": \"llama\",\n",
       "  \"num_attention_heads\": 32,\n",
       "  \"num_hidden_layers\": 32,\n",
       "  \"num_key_value_heads\": 8,\n",
       "  \"pretraining_tp\": 1,\n",
       "  \"rms_norm_eps\": 1e-05,\n",
       "  \"rope_scaling\": null,\n",
       "  \"rope_theta\": 500000.0,\n",
       "  \"tie_word_embeddings\": false,\n",
       "  \"torch_dtype\": \"bfloat16\",\n",
       "  \"transformers_version\": \"4.45.0.dev0\",\n",
       "  \"use_cache\": true,\n",
       "  \"vocab_size\": 128256\n",
       "}"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "llama3.config"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8289fc8-b8c4-4c50-8b6d-f2e2cc20806e",
   "metadata": {},
   "source": [
    "```\n",
    "class TransformerBlock(nn.Module):\n",
    "    def __init__(self, layer_id: int, args: ModelArgs):\n",
    "        super().__init__()\n",
    "        self.n_heads = args.n_heads\n",
    "        self.dim = args.dim\n",
    "        self.head_dim = args.dim // args.n_heads\n",
    "        self.attention = Attention(args)\n",
    "        self.feed_forward = FeedForward(\n",
    "            dim=args.dim,\n",
    "            hidden_dim=4 * args.dim,\n",
    "            multiple_of=args.multiple_of,\n",
    "            ffn_dim_multiplier=args.ffn_dim_multiplier,\n",
    "        )\n",
    "        self.layer_id = layer_id\n",
    "        self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)\n",
    "        self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)\n",
    "\n",
    "    def forward(\n",
    "        self,\n",
    "        x: torch.Tensor,\n",
    "        start_pos: int,\n",
    "        freqs_cis: torch.Tensor,\n",
    "        mask: Optional[torch.Tensor],\n",
    "    ):\n",
    "        h = x + self.attention(self.attention_norm(x), start_pos, freqs_cis, mask)\n",
    "        out = h + self.feed_forward(self.ffn_norm(h))\n",
    "        return out\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "7987c91d-d086-4c54-8e9d-06ce876ce80e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-09-16T02:52:03.923063Z",
     "iopub.status.busy": "2024-09-16T02:52:03.922848Z",
     "iopub.status.idle": "2024-09-16T02:52:04.051639Z",
     "shell.execute_reply": "2024-09-16T02:52:04.050867Z",
     "shell.execute_reply.started": "2024-09-16T02:52:03.923047Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([0.0537, 0.2090, 0.4492,  ..., 0.0859, 0.0437, 0.0292], device='cuda:0',\n",
       "       dtype=torch.bfloat16, requires_grad=True)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "llama3.model.layers[0].input_layernorm.weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b748fd11-e458-4130-92e0-db8ce48b8278",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
